import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 检查CUDA是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 定义3层神经网络
class ThreeLayerNet(nn.Module):
    def __init__(self, input_size, num_classes):
        super(ThreeLayerNet, self).__init__()
        self.input_size = input_size
        self.fc1 = nn.Linear(input_size, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, num_classes)
        self.relu = nn.ReLU()
        
    def forward(self, x):
        x = x.view(-1, self.input_size)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 加载数据集
def load_dataset(dataset_name, batch_size=64):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    
    if dataset_name == "FashionMNIST":
        train_set = torchvision.datasets.FashionMNIST(
            root='./data', 
            train=True, 
            download=True, 
            transform=transform
        )
        test_set = torchvision.datasets.FashionMNIST(
            root='./data', 
            train=False, 
            download=True, 
            transform=transform
        )
        input_size = 28*28
        num_classes = 10
    elif dataset_name == "CIFAR10":
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        train_set = torchvision.datasets.CIFAR10(
            root='./data', 
            train=True, 
            download=True, 
            transform=transform
        )
        test_set = torchvision.datasets.CIFAR10(
            root='./data', 
            train=False, 
            download=True, 
            transform=transform
        )
        input_size = 3*32*32
        num_classes = 10
    
    train_loader = torch.utils.data.DataLoader(
        train_set, 
        batch_size=batch_size, 
        shuffle=True
    )
    
    test_loader = torch.utils.data.DataLoader(
        test_set, 
        batch_size=batch_size, 
        shuffle=False
    )
    
    return train_loader, test_loader, input_size, num_classes

# 训练模型
def train_model(model, train_loader, test_loader, criterion, optimizer, num_epochs=10):
    train_losses = []
    train_accuracies = []
    test_losses = []
    test_accuracies = []
    
    model.train()
    
    for epoch in range(num_epochs):
        running_loss = 0.0
        correct = 0
        total = 0
        
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
        
        train_loss = running_loss / len(train_loader)
        train_accuracy = 100 * correct / total
        train_losses.append(train_loss)
        train_accuracies.append(train_accuracy)
        
        # 测试集评估
        test_loss, test_accuracy = evaluate_model(model, test_loader, criterion)
        test_losses.append(test_loss)
        test_accuracies.append(test_accuracy)
        
        print(f'Epoch [{epoch+1}/{num_epochs}], '
              f'Train Loss: {train_loss:.4f}, Train Acc: {train_accuracy:.2f}%, '
              f'Test Loss: {test_loss:.4f}, Test Acc: {test_accuracy:.2f}%')
    
    return train_losses, train_accuracies, test_losses, test_accuracies

# 评估模型
def evaluate_model(model, test_loader, criterion):
    model.eval()
    test_loss = 0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            test_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    
    test_loss = test_loss / len(test_loader)
    test_accuracy = 100 * correct / total
    
    return test_loss, test_accuracy

# 可视化训练过程
def visualize_training(train_losses, train_accuracies, test_losses, test_accuracies, dataset_name):
    plt.figure(figsize=(12, 5))
    
    # 损失曲线
    plt.subplot(1, 2, 1)
    plt.plot(train_losses, label='Train Loss')
    plt.plot(test_losses, label='Test Loss')
    plt.title(f'{dataset_name} - Training and Test Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    
    # 准确率曲线
    plt.subplot(1, 2, 2)
    plt.plot(train_accuracies, label='Train Accuracy')
    plt.plot(test_accuracies, label='Test Accuracy')
    plt.title(f'{dataset_name} - Training and Test Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.legend()
    
    plt.tight_layout()
    plt.show()

# 可视化预测结果
def visualize_predictions(model, test_loader, dataset_name):
    model.eval()
    dataiter = iter(test_loader)
    images, labels = next(dataiter)
    images, labels = images.to(device), labels.to(device)
    
    with torch.no_grad():
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
    
    # 显示图像和预测结果
    images = images.cpu()
    plt.figure(figsize=(10, 4))
    for i in range(6):
        plt.subplot(2, 3, i+1)
        if dataset_name == "FashionMNIST":
            img = images[i].numpy().squeeze()
            plt.imshow(img, cmap='gray')
        else:  # CIFAR10
            img = images[i].numpy().transpose((1, 2, 0))
            img = img * 0.5 + 0.5  # 反归一化
            plt.imshow(img)
        plt.title(f'Predicted: {predicted[i]}, Actual: {labels[i]}')
        plt.axis('off')
    plt.suptitle(f'{dataset_name} - 预测结果')
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    # 训练FashionMNIST
    print("=== 训练FashionMNIST ===")
    train_loader, test_loader, input_size, num_classes = load_dataset("FashionMNIST")
    model = ThreeLayerNet(input_size, num_classes).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    print("开始训练...")
    train_losses, train_accuracies, test_losses, test_accuracies = train_model(
        model, train_loader, test_loader, criterion, optimizer, num_epochs=15
    )
    
    # 可视化FashionMNIST结果
    visualize_training(train_losses, train_accuracies, test_losses, test_accuracies, "FashionMNIST")
    visualize_predictions(model, test_loader, "FashionMNIST")
    
    # 训练CIFAR10
    print("\n=== 训练CIFAR10 ===")
    train_loader, test_loader, input_size, num_classes = load_dataset("CIFAR10")
    model = ThreeLayerNet(input_size, num_classes).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    print("开始训练...")
    train_losses, train_accuracies, test_losses, test_accuracies = train_model(
        model, train_loader, test_loader, criterion, optimizer, num_epochs=15
    )
    
    # 可视化CIFAR10结果
    visualize_training(train_losses, train_accuracies, test_losses, test_accuracies, "CIFAR10")
    visualize_predictions(model, test_loader, "CIFAR10")